"""
MCP协议客户端实现（基于OpenAI SDK）
功能说明：
1. 支持多模型调用（通过OpenAI API）
2. 实现MCP消息循环机制
3. 自动处理响应流
4. 支持系统指令设置
"""

import openai
from typing import Dict, Any, Optional
from mcp_protocol import Message, Role

class MCPClient:
    """MCP协议客户端核心类"""
    
    def __init__(self, api_key: str, model: str = "gpt-4"):
        """
        初始化客户端
        Args:
            api_key: OpenAI API密钥
            model: 默认模型名称（支持gpt-4/gpt-3.5-turbo等）
        """
        self.client = openai.OpenAI(api_key=api_key)
        self.model = model
        self.system_prompt = None  # 系统级指令

    def set_system_prompt(self, prompt: str):
        """设置系统级指令"""
        self.system_prompt = prompt

    def _format_mcp_message(self, message: Dict) -> Dict:
        """将MCP消息格式转换为OpenAI格式"""
        formatted = {
            "role": message.get("role", "user"),
            "content": message.get("content", "")
        }
        if self.system_prompt:
            formatted["role"] = "system"
            formatted["content"] = self.system_prompt
        return formatted

    def send_message(self, message: Dict, stream: bool = False) -> Any:
        """
        发送MCP消息并获取响应
        Args:
            message: MCP消息对象
            stream: 是否启用流式响应
        Returns:
            OpenAI响应对象
        """
        try:
            formatted_messages = [self._format_mcp_message(msg) for msg in message.get("messages", [])]
            
            response = self.client.chat.completions.create(
                model=self.model,
                messages=formatted_messages,
                stream=stream,
                temperature=message.get("temperature", 0.7),
                max_tokens=message.get("max_tokens", 1024)
            )
            return response
        except Exception as e:
            print(f"API调用错误: {str(e)}")
            return None

    def process_response(self, response: Any):
        """处理模型响应"""
        if response and response.choices:
            return response.choices[0].message.content
        return "未获取到有效响应"

# 示例用法
if __name__ == "__main__":
    # 初始化客户端
    client = MCPClient(api_key="your-api-key")
    client.set_system_prompt("你是一个专业的科技助手，回答要简洁准确")
    
    # 发送MCP消息
    message = {
        "messages": [
            {"role": "user", "content": "什么是MCP协议？"},
            {"role": "assistant", "content": "MCP全称Model Context Protocol"}
        ],
        "temperature": 0.5,
        "max_tokens": 500
    }
    
    # 获取并处理响应
    response = client.send_message(message)
    print("模型响应:", client.process_response(response))
